"""WebSocket connection manager."""

from __future__ import annotations

from typing import Dict, Set

from fastapi import WebSocket


class ConnectionManager:
    def __init__(self) -> None:
        self.active_connections: Set[WebSocket] = set()
        self.analysis_connections: Dict[int, Set[WebSocket]] = {}

    async def connect(self, websocket: WebSocket) -> None:
        await websocket.accept()
        self.active_connections.add(websocket)

    def disconnect(self, websocket: WebSocket) -> None:
        self.active_connections.discard(websocket)
        for connections in self.analysis_connections.values():
            connections.discard(websocket)

    async def send_personal_message(self, event: str, payload: str, websocket: WebSocket) -> None:
        await websocket.send_json({"event": event, "payload": payload})

    async def broadcast(self, event: str, payload: str) -> None:
        for connection in list(self.active_connections):
            await connection.send_json({"event": event, "payload": payload})

    async def connect_analysis(self, project_id: int, websocket: WebSocket) -> None:
        await websocket.accept()
        self.analysis_connections.setdefault(project_id, set()).add(websocket)

    async def broadcast_analysis(
        self, project_id: int, event: str, payload: dict[str, str | float | int]
    ) -> None:
        for connection in list(self.analysis_connections.get(project_id, set())):
            await connection.send_json(
                {"project_id": project_id, "event": event, "payload": payload}
            )


manager = ConnectionManager()
